package z1util

import (
	"context"

	"gorm.io/gorm"
)

// update the context of gorm.DB
func UpdateDBContext(key string, val any, db *gorm.DB, replace bool) (db2 *gorm.DB) {
	ctx := db.Statement.Context

	if replace {
		ctx = context.Background()
	}

	// 当replace==true时，全部的context.WithValue替换为最后设置的一个；否则context.WithValue存在就更新，不存在就追加
	db2 = db.WithContext(context.WithValue(
		ctx,
		key,
		val,
	))

	return
}
